import functools
import asyncio
import logging
from typing import TypeVar, Callable, Any, Optional
import time
import tiktoken

logger = logging.getLogger(__name__)

T = TypeVar("T")

def get_token_count(text: str) -> int:
    """
    计算文本的token数量
    
    Args:
        text: 需要计算token数的文本
        
    Returns:
        token数量
    """
    try:
        # 使用cl100k_base编码器(GPT-4使用的编码器)
        encoding = tiktoken.get_encoding("cl100k_base")
        return len(encoding.encode(text))
    except Exception as e:
        logger.error(f"计算token数量时发生错误: {str(e)}")
        # 如果发生错误,使用简单的估算方法(每4个字符约等于1个token)
        return len(text) // 4

def async_retry(
    retries: int = 3,
    delay: float = 1.0,
    backoff: float = 2.0,
    exceptions: tuple = (Exception,),
    on_retry: Optional[Callable[[Exception, int], None]] = None
):
    """
    异步函数重试装饰器
    
    Args:
        retries: 最大重试次数
        delay: 初始延迟时间(秒)
        backoff: 延迟时间的增长因子
        exceptions: 需要重试的异常类型
        on_retry: 重试时的回调函数
    """
    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        @functools.wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> Any:
            last_exception = None
            current_delay = delay

            for attempt in range(retries + 1):
                try:
                    return await func(*args, **kwargs)
                except exceptions as e:
                    last_exception = e
                    if attempt == retries:
                        logger.error(
                            f"最后一次重试失败 {func.__name__}: {str(e)}"
                        )
                        raise

                    if on_retry:
                        on_retry(e, attempt + 1)

                    logger.warning(
                        f"{func.__name__} 第{attempt + 1}次重试, 错误: {str(e)}, "
                        f"等待 {current_delay} 秒后重试"
                    )
                    await asyncio.sleep(current_delay)
                    current_delay *= backoff

            raise last_exception

        return wrapper
    return decorator

def sync_retry(
    retries: int = 3,
    delay: float = 1.0,
    backoff: float = 2.0,
    exceptions: tuple = (Exception,),
    on_retry: Optional[Callable[[Exception, int], None]] = None
):
    """
    同步函数重试装饰器
    
    Args:
        retries: 最大重试次数
        delay: 初始延迟时间(秒)
        backoff: 延迟时间的增长因子
        exceptions: 需要重试的异常类型
        on_retry: 重试时的回调函数
    """
    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            last_exception = None
            current_delay = delay

            for attempt in range(retries + 1):
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    last_exception = e
                    if attempt == retries:
                        logger.error(
                            f"最后一次重试失败 {func.__name__}: {str(e)}"
                        )
                        raise

                    if on_retry:
                        on_retry(e, attempt + 1)

                    logger.warning(
                        f"{func.__name__} 第{attempt + 1}次重试, 错误: {str(e)}, "
                        f"等待 {current_delay} 秒后重试"
                    )
                    time.sleep(current_delay)
                    current_delay *= backoff

            raise last_exception

        return wrapper
    return decorator 